#pragma once
#include "Ray.h"
#include "RayColumn.h"
#include "Intersection.h"
#include "Scene.h"
#include "Heightmap.h"
#include <stdio.h>
#include <float.h>

#define EPSILON 0.001

using namespace std;

template <class CellType>
class KdTreeCellScene;

template <class CellType>
class KdTreeCellSceneNode {
private:
    friend KdTreeCellScene<CellType>;
    KdTreeCellSceneNode<CellType>* negChild;
    KdTreeCellSceneNode<CellType>* posChild;
    CellType* leaf; //null if not leaf node
    int split;
    float minHeight, maxHeight;
    bool axis; //true = y axis

    void findSplit(Ray& ray, KdTreeCellSceneNode<CellType>*& nearChild, KdTreeCellSceneNode<CellType>*& farChild, float& splitT) const;
public:
    //Traces a ray into the scene. Returns the intersection.
    Intersection traceRay(Ray& ray, float minT, float maxT, float minZ) const;

    //Traces a shadow ray. Returns true iff there is a collision.
    bool traceShadowRay(Ray& ray, float minT, float maxT, float minZ, float maxZ) const;

    //Traces a SlopeRayColumn beginning at the start-th ray.
    int traceSlopeRayColumn(Intersection* results, SlopeRayColumn& slopeRayColumn, float minT, float maxT, int start) const;
    
    static KdTreeCellSceneNode<CellType>* createFromHeightmap(Heightmap& heightmap, int minX, int maxX, int minY, int maxY);
};

template <class CellType>
void KdTreeCellSceneNode<CellType>::findSplit(Ray& ray, KdTreeCellSceneNode<CellType>*& nearChild, KdTreeCellSceneNode<CellType>*& farChild, float& splitT) const {
    if (axis) {
        splitT = ray.timeToY(split);
        if (ray.dir.y > 0) {
            nearChild = negChild;
            farChild = posChild;
        } else {
            nearChild = posChild;
            farChild = negChild;
        }
    } else {
        splitT = ray.timeToX(split);
        if (ray.dir.x > 0) {
            nearChild = negChild;
            farChild = posChild;
        } else {
            nearChild = posChild;
            farChild = negChild;
        }
    }
}

template <class CellType>
Intersection KdTreeCellSceneNode<CellType>::traceRay(Ray& ray, float minT, float maxT, float minZ) const {
    if (minZ > maxHeight) return Intersection();

    if (leaf) {
        return leaf->traceRay(ray);
    }

    KdTreeCellSceneNode<CellType>* nearChild;
    KdTreeCellSceneNode<CellType>* farChild;
    float splitT;

    findSplit(ray, nearChild, farChild, splitT);

    Intersection result;

    if (minT <= splitT) {
        float negMaxT = maxT;
        float negMinZ = minZ;
        if (splitT < maxT) {
            negMaxT = splitT;
            if (ray.dir.z < 0) negMinZ = ray.origin.z + ray.dir.z * negMaxT;
        }
        result = nearChild->traceRay(ray, minT, negMaxT, negMinZ);
    }

    if (result.exists()) return result;

    if (splitT <= maxT) {
        float posMinT = minT;
        float posMinZ = minZ;
        if (splitT > minT) {
            posMinT = splitT;
            if (ray.dir.z > 0) posMinZ = ray.origin.z + ray.dir.z * posMinT;
        }
        result = farChild->traceRay(ray, posMinT, maxT, posMinZ);
    }

    return result;
}

template <class CellType>
bool KdTreeCellSceneNode<CellType>::traceShadowRay(Ray& ray, float minT, float maxT, float minZ, float maxZ) const {
    if (minZ > maxHeight) return false;
    if (maxZ < minHeight + EPSILON) return true;

    if (leaf) {
        return leaf->traceShadowRay(ray);
    }

    KdTreeCellSceneNode<CellType>* nearChild;
    KdTreeCellSceneNode<CellType>* farChild;
    float splitT;

    findSplit(ray, nearChild, farChild, splitT);

    if (minT <= splitT) {
        float negMaxT = maxT;
        float negMinZ = minZ;
        float negMaxZ = maxZ;
        if (splitT < maxT) {
            negMaxT = splitT;
            if (ray.dir.z < 0) negMinZ = ray.origin.z + ray.dir.z * negMaxT;
            else negMaxZ = ray.origin.z + ray.dir.z * negMaxT;
        }
        if (nearChild->traceShadowRay(ray, minT, negMaxT, negMinZ, negMaxZ)) return true;
    }

    if (splitT <= maxT) {
        float posMinT = minT;
        float posMinZ = minZ;
        float posMaxZ = maxZ;
        if (splitT > minT) {
            posMinT = splitT;
            if (ray.dir.z > 0) posMinZ = ray.origin.z + ray.dir.z * posMinT;
            else posMaxZ = ray.origin.z + ray.dir.z * posMinT;
        }
        if (farChild->traceShadowRay(ray, posMinT, maxT, posMinZ, posMaxZ)) return true;
    }

    return false;
}

template <class CellType>
int KdTreeCellSceneNode<CellType>::traceSlopeRayColumn(Intersection* results, SlopeRayColumn& slopeRayColumn, float minT, float maxT, int start) const {
    float minZ = slopeRayColumn.origin.z + minT * slopeRayColumn.dirZ[start];
    float maxZ = slopeRayColumn.origin.z + maxT * slopeRayColumn.dirZ[start];
    minZ = min(minZ, maxZ);
    if (minZ > maxHeight) return start;

    if (leaf) {
        for (; start < slopeRayColumn.numRays; start++) {
            Intersection result = leaf->traceRay(slopeRayColumn.getRay(start));
            if (result.exists()) {
                results[start] = result;
            } else {
                break;
            }
        }
        return start;
    }

    KdTreeCellSceneNode<CellType>* nearChild;
    KdTreeCellSceneNode<CellType>* farChild;
    float splitT;

    findSplit(slopeRayColumn.flatRay, nearChild, farChild, splitT);

    if (minT <= splitT) {
        float negMaxT = maxT;
        if (splitT < maxT) {
            negMaxT = splitT;
        }
        start = nearChild->traceSlopeRayColumn(results, slopeRayColumn, minT, negMaxT, start);
    }

    if (splitT <= maxT) {
        float posMinT = minT;
        if (splitT > minT) {
            posMinT = splitT;
        }
        start = farChild->traceSlopeRayColumn(results, slopeRayColumn, posMinT, maxT, start);
    }
    
    return start;
}

template <class CellType>
KdTreeCellSceneNode<CellType>* KdTreeCellSceneNode<CellType>::createFromHeightmap(Heightmap& heightmap, int minX, int maxX, int minY, int maxY) {
    KdTreeCellSceneNode<CellType>* result = new KdTreeCellSceneNode<CellType>();
    if (maxX == minX + 1 && maxY == minY + 1) {
        //leaf node
        result->negChild = 0;
        result->posChild = 0;
        result->leaf = new CellType();
        result->leaf->setFromHeightmap(heightmap, minX, minY);
        result->minHeight = result->leaf->getMinHeight();
        result->maxHeight = result->leaf->getMaxHeight();
        return result;
    }

    //non-leaf nodes
    result->leaf = 0;

    /**
     * Choose an axis and a split axis.
     * Use a surface area heuristic: minimize surface area.
     * Top surface area is constant with split choice, so only worry about sides.
     */

    float xScore = FLT_MAX;
    float yScore = FLT_MAX;
    int xSplit = 0;
    int ySplit = 0;

    if (maxX > minX + 1) {
        //placeholder: even split
        xSplit = (maxX + minX) / 2;
        xScore = float(maxY - minY);

        //heightmap.getOptSplitX(xScore, xSplit, minX, maxX + 1, minY, maxY + 1);
    }

    if (maxY > minY + 1) {
        //placeholder: even split
        ySplit = (minY + maxY) / 2;
        yScore = float(maxX - minX);

        //heightmap.getOptSplitY(yScore, ySplit, minX, maxX + 1, minY, maxY + 1);
    }

    if (xScore < yScore) {
        result->negChild = createFromHeightmap(heightmap, minX, xSplit, minY, maxY);
        result->posChild = createFromHeightmap(heightmap, xSplit, maxX, minY, maxY);
        result->split = xSplit;
        result->axis = false;
    } else {
        result->negChild = createFromHeightmap(heightmap, minX, maxX, minY, ySplit);
        result->posChild = createFromHeightmap(heightmap, minX, maxX, ySplit, maxY);
        result->split = ySplit;
        result->axis = true;
    }

    result->maxHeight = max(result->negChild->maxHeight, result->posChild->maxHeight);
    result->minHeight = min(result->negChild->minHeight, result->posChild->minHeight);

    return result;
}


template <class CellType>
class KdTreeCellScene:
    public Scene {
public:
    private:
    CellType* cells;

    int maxX, maxY;

    KdTreeCellSceneNode<CellType>* root;

    Light* userLight;

    Color3f shadeHit(Intersection& hit) const;
public:
    KdTreeCellScene();

    //Readies this scene for ray tracing.
    void ready();

    //Assigns a heightmap to this scene.
    void assignHeightmap(Heightmap& heightmap);

    //Assigns a user-controlled light.
    void assignUserLight(Light* light);

    //Modify the user light.
    void modifyUserLight(float x, float y);

    //Traces a ray into the scene. Returns the intersection.
    Color3f traceRay(Ray& ray) const;

    //Traces a shadow ray. Returns true iff there is a collision.
    bool traceShadowRay(Ray& ray) const;

    bool traceDelayedShadowRay(Ray& ray, float delay) const;

    //Traces a SlopeRayColumn into the scene. Default iterates over all rays in the column.
    void traceSlopeRayColumn(Color3f* results, SlopeRayColumn& slopeRayColumn) const;
};

template <class CellType>
KdTreeCellScene<CellType>::KdTreeCellScene() {
    userLight = 0;
}

template <class CellType>
void KdTreeCellScene<CellType>::ready() {
    
}

template <class CellType>
void KdTreeCellScene<CellType>::assignHeightmap(Heightmap& heightmap) {
    maxX = heightmap.width() - 1;
    maxY = heightmap.height() - 1;
    root = KdTreeCellSceneNode<CellType>::createFromHeightmap(heightmap, 0, maxX, 0, maxY);
}

template <class CellType>
void KdTreeCellScene<CellType>::assignUserLight(Light* light) {
    userLight = light;
}

template <class CellType>
void KdTreeCellScene<CellType>::modifyUserLight(float x, float y) {
    if (userLight) userLight->userModify(x, y);
}

template <class CellType>
Color3f KdTreeCellScene<CellType>::shadeHit(Intersection& hit) const {
    Color3f result;

    if (hit.exists()) {
        if (userLight) {
            Lighting lighting = userLight->getLightingFrom(hit.pos);
            result = Color3f(0.125, 0.125, 0.125);
            float obliquity = hit.n * lighting.dir;
            float obliquity_s = hit.n_s * lighting.dir;
            if (obliquity_s > 0.0) {
                if (obliquity > 0.0) {
                    Ray shadowRay = Ray(hit.pos, lighting.dir);
                    if (!traceShadowRay(shadowRay)) {
                        result += lighting.color * obliquity_s;
                    }
                } else {
                    Ray shadowRay = Ray(hit.pos + lighting.dir, lighting.dir);
                    float minT, maxT;

                    if (shadowRay.timeInRect(minT, maxT, 0, maxX, 0, maxY)) {
                        float minZ = shadowRay.origin.z + minT * shadowRay.dir.z;
                        float maxZ = shadowRay.origin.z + maxT * shadowRay.dir.z;
                        minZ = min(minZ, maxZ);

                        Intersection hit = root->traceRay(shadowRay, minT, maxT, minZ);
                        if (!hit.exists()) result += lighting.color * obliquity_s;
                    } else {
                        result += lighting.color * obliquity_s;
                    }
                }
            }
        } else {
            result = Color3f(hit.n_s.x, hit.n_s.y, hit.n_s.z);
        }
    } else {
        result = Color3f(0, 0, 0);
    }
    
    return result;
}

template <class CellType>
Color3f KdTreeCellScene<CellType>::traceRay(Ray& ray) const {
    //hits the root node square?
    float minT, maxT;

    if (!ray.timeInRect(minT, maxT, 0, maxX, 0, maxY)) return Color3f(0, 0, 0);

    float minZ = ray.origin.z + minT * ray.dir.z;
    float maxZ = ray.origin.z + maxT * ray.dir.z;
    minZ = min(minZ, maxZ);

    Intersection hit = root->traceRay(ray, minT, maxT, minZ);

    return shadeHit(hit);
}

template <class CellType>
bool KdTreeCellScene<CellType>::traceShadowRay(Ray& ray) const {
    //hits the root node square?
    float minT, maxT;

    if (!ray.timeInRect(minT, maxT, 0, maxX, 0, maxY)) return false;

    float minZ = ray.origin.z + minT * ray.dir.z;
    float maxZ = ray.origin.z + maxT * ray.dir.z;
    if (minZ > maxZ) swap(minZ, maxZ);
    return root->traceShadowRay(ray, minT, maxT, minZ, maxZ);
}

template <class CellType>
void KdTreeCellScene<CellType>::traceSlopeRayColumn(Color3f* results, SlopeRayColumn& slopeRayColumn) const {
    Intersection* hits = new Intersection[slopeRayColumn.numRays];

    for (int i = 0; i < slopeRayColumn.numRays; i++) hits[i] = Intersection();

    float minT, maxT;
    if (!slopeRayColumn.flatRay.timeInRect(minT, maxT, 0, maxX, 0, maxY)) {
        for (int i = 0; i < slopeRayColumn.numRays; i++) results[i] = Color3f(0, 0, 0);
        return;
    }

    root->traceSlopeRayColumn(hits, slopeRayColumn, minT, maxT, 0);

    for (int i = 0; i < slopeRayColumn.numRays; i++) results[i] = shadeHit(hits[i]);

    delete[] hits;
}
